from imagedownloader import yandexapi
from CLIP import clip_score as clip
from ipadapter import IpAdapter as ip
from ip2p import edit_cli as ip2p
from gpt import GPT
import sys
import os
sys.path.append('MAG-Edit/code_tr')
from edit import edit_image
import logging
import shutil
import re
import maskgen


if __name__ == '__main__':
    # 控制整体的生成图片顺序
    progress = 1

    # 设置日志级别: INFO
    logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                        level=logging.INFO)

    # 输入文本
    logging.info("###################  开始执行  ###################")
    input_text = "Husky wearing a straw hat on the snow"
    logging.info("Input_text: %s", input_text)

    # 保存图片的目录
    img_dir = "./images"
    logging.info("############  检索图片将保存在：%s  ###############", img_dir)

    # 原文本检索
    logging.info("############  开始检索原文本  ###############")
    yandexapi.search_image(input_text, img_dir, 30, progress)
    logging.info("############  检索原文本结束  ###############")

    # 调用GPT获得丰富检索词
    logging.info("############  开始调用GPT获得语义扩展  ###############")
    query = GPT.get_search_keywords(input_text)
    logging.info("############  语义扩展结束  ###############")

    # 执行检索策略
    logging.info("############  开始执行检索策略  ###############")
    for i in range(len(query)):
        yandexapi.search_image(query[i], img_dir, 5, 2)
    logging.info("############  检索策略结束  ###############")

    # 筛选器模块
    # CLIP评分模块
    logging.info("############  开始执行CLIP评分模块  ###############")
    prefix = "a photo of "
    clip_text = prefix + input_text
    pattern = re.compile(r'^image_(\d+)_\d+$', re.IGNORECASE)
    count = 0
    scores = {}
    for image in os.listdir(img_dir):
        basename, ext = os.path.splitext(image)
        match = pattern.match(basename)
        if int(match.group(1)) != progress:
            continue
        image_path = os.path.join(img_dir, image)
        try:
            score = clip.get_clip_score(image_path, clip_text)
            count += 1
        except Exception as e:
            logging.error(f"计算图片{image}的CLIP分数时出错：{e}")
            continue
        scores[image] = score
        logging.info(f"{image} CLIP 分数为：{score}")
    sort_images = sorted(scores, key=scores.get, reverse=True)[:count]
    logging.info("CLIP 分数排序：")
    print(sort_images)
    clip_top_image = sort_images[0]
    logging.info(f"CLIP 评分最高的图片: {clip_top_image}")
    logging.info("############  CLIP评分模块结束  ###############")

    # LLM筛选器
    logging.info("############  开始执行LLM筛选器  ###############")
    for image in sort_images:
        image_path = os.path.join(img_dir, image)
        if GPT.filter_image(image_path, input_text):
            clip_top_image = image
            print(f"筛选图片，最终编辑图片: {clip_top_image}")
            break
    logging.info("############  LLM筛选器结束  ###############")

    # 存储图片
    logging.info("############  编辑开始  ###############")
    backups_origin_dir = '/root/data/origins'  # 原始图片
    backups_output_dir = '/root/data/outputs'  # 生成图片
    # 创建文件夹
    if not os.path.exists(backups_origin_dir):
        os.makedirs(backups_origin_dir)
    if not os.path.exists(backups_output_dir):
        os.makedirs(backups_output_dir)
    source_path = os.path.join(img_dir, clip_top_image)
    name, ext = os.path.splitext(clip_top_image)
    backups_origin_dest_path = os.path.join(backups_origin_dir, f'{progress}' + ext)
    logging.info(f"=======>  原始图片存放在： {backups_origin_dest_path}  <=======")
    backups_output_dest_path = os.path.join(backups_output_dir, f'{progress}' + ext)
    logging.info(f"=======>  即将编辑这张图片： {backups_output_dest_path}  <=======")
    shutil.copy(source_path, backups_origin_dest_path)
    shutil.copy(source_path, backups_output_dest_path)

    # 调用大模型生成编辑指令
    logging.info("############  开始调用GPT获得编辑指令  ###############")
    instructions = GPT.get_diff_edit_instruction(backups_output_dest_path, input_text)

    """
    instructions按照如下格式组织：
    [
    {"Title": "", "Edit": "", "Before Object":"", "Object"},
    ......
    ]
    """
    logging.info("############  差异编辑开始  ###############")
    for instruction in instructions:
        # 生成掩码图片
        mask_img_path = maskgen.mask_generate(backups_output_dest_path, instruction['Before Object'], progress)
        # 差异编辑
        try:
            source_prompt = instruction['Title']  # 使用原始输入文本作为源提示词
            target_prompt = instruction['Edit']  # 使用编辑指令作为目标提示词
            target_word = [instruction['Object']]  # 目标对象

            # 调用edit_image方法
            images, x_t = edit_image(
                source_prompt=source_prompt,
                target_prompt=target_prompt,
                target_word=target_word,
                img_path=backups_output_dest_path,
                mask_path=mask_img_path,
                result_dir=f"edit_result_{progress}",
                max_iteration=15,
                scale=2.5
            )

            logging.info(f"成功编辑图片，指令：{instruction['Edit']}")

        except Exception as e:
            logging.error("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
            logging.error(f"编辑图片时出错：{e}")
            logging.error("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
            continue

    logging.info("############  执行结束  ###############")
